"""
动态元数据提取器
使用LLM智能提取新闻的元数据
"""

import json
from typing import Dict, List, Any

class MetadataExtractor:
    """动态元数据提取器"""
    
    def __init__(self, llm):
        self.llm = llm
    
    def extract_metadata(self, news_data: Dict[str, Any]) -> Dict[str, Any]:
        """从新闻数据中动态提取元数据"""
        prompt = f"""
请从以下新闻中提取关键元数据：

新闻标题：{news_data.get("title", "")}
新闻内容：{news_data.get("content", "")}

请提取以下信息，并以JSON格式返回：
1. tokens: 代币符号列表
2. institutions: 机构名称列表
3. projects: 项目名称列表
4. event_types: 事件类型列表
5. keywords: 关键词列表
6. categories: 分类列表

返回格式：
{{
    "tokens": [],
    "institutions": [],
    "projects": [],
    "event_types": [],
    "keywords": [],
    "categories": []
}}
"""
        
        response = self.llm.invoke(prompt)
        return self._parse_response(response.content)
    
    def _parse_response(self, response: str) -> Dict[str, Any]:
        """解析LLM响应"""
        try:
            if response.strip().startswith('{'):
                return json.loads(response.strip())
            
            import re
            json_match = re.search(r'\{.*\}', response, re.DOTALL)
            if json_match:
                return json.loads(json_match.group())
            
            return self._get_default_metadata()
        except:
            return self._get_default_metadata()
    
    def _get_default_metadata(self) -> Dict[str, Any]:
        """获取默认元数据"""
        return {
            "tokens": [],
            "institutions": [],
            "projects": [],
            "event_types": [],
            "keywords": [],
            "categories": []
        }
    
    def generate_queries(self, metadata: Dict[str, Any], news_data: Dict[str, Any]) -> List[Dict[str, str]]:
        """基于元数据生成检索查询"""
        prompt = f"""
基于以下新闻和元数据，生成3-5个用于知识库检索的问题：

新闻标题：{news_data.get("title", "")}
新闻内容：{news_data.get("content", "")}

元数据：
- 代币：{metadata.get('tokens', [])}
- 机构：{metadata.get('institutions', [])}
- 项目：{metadata.get('projects', [])}
- 事件类型：{metadata.get('event_types', [])}
- 关键词：{metadata.get('keywords', [])}

返回格式：
[
    {{"query": "问题1", "type": "historical"}},
    {{"query": "问题2", "type": "influence"}},
    {{"query": "问题3", "type": "risk"}}
]
"""
        
        response = self.llm.invoke(prompt)
        return self._parse_queries(response.content)
    
    def _parse_queries(self, response: str) -> List[Dict[str, str]]:
        """解析查询响应"""
        try:
            if response.strip().startswith('['):
                return json.loads(response.strip())
            
            import re
            json_match = re.search(r'\[.*\]', response, re.DOTALL)
            if json_match:
                return json.loads(json_match.group())
            
            return self._get_default_queries()
        except:
            return self._get_default_queries()
    
    def _get_default_queries(self) -> List[Dict[str, str]]:
        """获取默认查询"""
        return [
            {"query": "相关项目的融资历史", "type": "historical"},
            {"query": "投资机构的历史表现", "type": "influence"},
            {"query": "市场风险和机会", "type": "risk"}
        ]